#!/usr/bin/env python
# -*- coding: utf-8 -*-
import sys
import os
import yaml
import math
import threading
import cv2
import time
from queue import Queue
import numpy as np
from typing import Union
import argparse
import platform
from spirems import Subscriber, Publisher, cvimg2sms, sms2cvimg, def_msg, QoS, BaseNode, get_extra_args
from spirems.mod_helper import download_model
from spirecv.algorithm.utils import calc_fov, calc_los_pos, draw_bbox


cropping = False
ref_point = []


def box_selection(event, x, y, flags, param):
    global ref_point, cropping, img, clone
    if event == cv2.EVENT_LBUTTONDOWN:
        ref_point = [(x, y)]
        cropping = True
    elif event == cv2.EVENT_MOUSEMOVE:
        if cropping:
            img_copy = clone.copy()
            cv2.rectangle(img_copy, ref_point[0], (x, y), (0, 0, 255), 2)
            cv2.imshow("trackwin", img_copy)
    elif event == cv2.EVENT_LBUTTONUP:
        ref_point.append((x, y))
        cropping = False
        cv2.rectangle(img, ref_point[0], (x, y), (0, 0, 255), 2)
        cv2.imshow("trackwin", img)


class ClickTrackControlNode(threading.Thread, BaseNode):
    def __init__(
        self,
        job_name: str,
        ip: str = '127.0.0.1',
        port: int = 9094,
        param_dict_or_file: Union[dict, str] = None,
        sms_shutdown: bool = True,
        **kwargs
    ):
        threading.Thread.__init__(self)
        sms_shutdown = True if sms_shutdown in ['True', 'true', '1', True] else False
        BaseNode.__init__(
            self,
            self.__class__.__name__,
            job_name,
            ip=ip,
            port=port,
            param_dict_or_file=param_dict_or_file,
            sms_shutdown=sms_shutdown,
            **kwargs
        )
        self.specified_input_topic = self.get_param("specified_input_topic", "")
        self.specified_output_topic = self.get_param("specified_output_topic", "")
        self.realtime_det = self.get_param("realtime_det", False)
        self.remote_ip = self.get_param("remote_ip", "127.0.0.1")
        self.remote_port = self.get_param("remote_port", 9094)
        self.use_shm = self.get_param("use_shm", -1)
        self.params_help()

        self.b_use_shm = False
        if self.use_shm == 1 or (self.use_shm == -1 and platform.system() == 'Linux'):
            self.b_use_shm = True

        input_url = '/' + job_name + '/sensor/image_raw'
        if len(self.specified_input_topic) > 0:
            input_url = specified_input_topic

        output_url = '/' + job_name + '/track/results'
        if len(self.specified_output_topic) > 0:
            output_url = specified_output_topic

        calib_url = '/' + job_name + '/sensor/calibration_info'

        self.job_queue = Queue()
        self.queue_pool.append(self.job_queue)
        self.detect_or_track = 0

        self.calib_width, self.calib_height = -1, -1
        self.camera_matrix = [712.12, 0,645.23, 0, 705.87, 327.34, 0, 0, 1]
        self.camera_matrix = np.array(self.camera_matrix).reshape(3, 3)
        self.distortion = [0.0, 0.0, 0.0, 0.0, 0.0]
        self.distortion = np.array(self.distortion)

        self._image_reader = Subscriber(
            input_url, 'std_msgs::Null', self.image_callback,
            ip=ip, port=port
        )
        self._calibration_reader = Subscriber(
            calib_url, 'sensor_msgs::CameraCalibration', self.calibration_callback,
            ip=ip, port=port, qos=QoS.Reliability
        )
        self._result_writer = Publisher(
            output_url, 'spirecv_msgs::2DTargets',
            ip=self.remote_ip, port=self.remote_port, qos=QoS.Reliability
        )
        self._show_writer = Publisher(
            '/' + job_name + '/track/image_results', 'memory_msgs::RawImage' if self.b_use_shm else 'sensor_msgs::CompressedImage',
            ip=ip, port=port
        )

        self._detect_writer = Publisher(
            '/click_detect/sensor/image_raw', 'memory_msgs::RawImage' if self.b_use_shm else 'sensor_msgs::CompressedImage',
            ip=ip, port=port
        )
        self._track_writer = Publisher(
            '/click_track/sensor/image_raw', 'memory_msgs::RawImage' if self.b_use_shm else 'sensor_msgs::CompressedImage',
            ip=ip, port=port
        )
        self._detect_calib_writer = Publisher(
            '/click_detect/sensor/calibration_info', 'sensor_msgs::CameraCalibration',
            ip=ip, port=port
        )
        self._track_calib_writer = Publisher(
            '/click_track/sensor/calibration_info', 'sensor_msgs::CameraCalibration',
            ip=ip, port=port
        )
        self._detect_result_reader = Subscriber(
            '/click_detect/detector/image_results', 'std_msgs::Null', self.detect_result_callback,
            ip=ip, port=port
        )
        self._track_result_reader = Subscriber(
            '/click_track/track/image_results', 'std_msgs::Null', self.track_result_callback,
            ip=ip, port=port
        )
        self._track_info_writer = Publisher(
            '/click_track/info', 'spirecv_msgs::ClickInfo',
            ip=ip, port=port
        )

        self.start()

    def release(self):
        BaseNode.release(self)
        self._image_reader.kill()
        self._calibration_reader.kill()
        self._result_writer.kill()
        self._show_writer.kill()
        self._detect_writer.kill()
        self._track_writer.kill()
        self._detect_result_reader.kill()
        self._track_result_reader.kill()

    def image_callback(self, msg):
        if self.detect_or_track == 0:
            self._detect_writer.publish(msg)
        elif self.detect_or_track == 1:
            self._track_writer.publish(msg)

    def detect_result_callback(self, msg):
        if self.detect_or_track == 0:
            img = sms2cvimg(msg)
            self.job_queue.put({'msg': msg, 'img': img})
    
    def track_result_callback(self, msg):
        if self.detect_or_track == 1:
            img = sms2cvimg(msg)
            self.job_queue.put({'msg': msg, 'img': img})

    def calibration_callback(self, msg):
        if self.detect_or_track == 0:
            self._detect_calib_writer.publish(msg)
        elif self.detect_or_track == 1:
            self._track_calib_writer.publish(msg)

        self.calib_width = msg['width']
        self.calib_height = msg['height']

        self.camera_matrix = np.array(msg['K']).reshape(3,3)
        self.distortion = np.array(msg['D'])

    def run(self):
        global img, clone, cropping, ref_point
        init_rect = [0, 0, 0, 0]  # xywh
        count = 0

        while self.is_running():
            t1 = time.time()
            msg_dict = self.job_queue.get(block=True)
            if msg_dict is None:
                break

            msg, img = msg_dict['msg'], msg_dict['img']

            clone = img.copy()
            file_name = msg['file_name'] if 'file_name' in msg else ''
            img = draw_bbox(img, msg)
            self._result_writer.publish(msg['spirecv_msgs::2DTargets'])
            show_msg = self._show_writer.cvimg2sms_mem(img)
            self._show_writer.publish(show_msg)

            cv2.namedWindow("trackwin")
            cv2.setMouseCallback("trackwin", box_selection)
            cv2.imshow("trackwin", img)
            key = cv2.waitKey(5)
            if key == 27:
                self.detect_or_track = 0

            if len(ref_point) == 2:
                if self.detect_or_track == 0:
                    print(ref_point)
                    init_rect = [ref_point[0][0], ref_point[0][1], ref_point[1][0] - ref_point[0][0], ref_point[1][1] - ref_point[0][1]]
                    if init_rect[2] >= 5 and init_rect[3] >= 5:
                        info_msg = def_msg("spirecv_msgs::ClickInfo")
                        info_msg["bbox"] = init_rect
                        self._track_info_writer.publish(info_msg)
                        self.detect_or_track = 1
                        self.job_queue.queue.clear()
                    else:
                        cx, cy = ref_point[1]
                        found = False
                        box_sel = []
                        min_dis = 1e6
                        for obj in msg['spirecv_msgs::2DTargets']['targets']:
                            if obj['bbox'][0] < cx < obj['bbox'][0] + obj['bbox'][2] and obj['bbox'][1] < cy < obj['bbox'][1] + obj['bbox'][3]:
                                ox = obj['bbox'][0] + obj['bbox'][2] / 2.
                                oy = obj['bbox'][1] + obj['bbox'][3] / 2.
                                dis = math.sqrt((cx - ox) ** 2 + (cy - oy) ** 2)
                                if dis < min_dis:
                                    min_dis = dis
                                    found = True
                                    box_sel = obj['bbox']
                        if found:
                            info_msg = def_msg("spirecv_msgs::ClickInfo")
                            info_msg["bbox"] = box_sel
                            self._track_info_writer.publish(info_msg)
                            self.detect_or_track = 1
                            self.job_queue.queue.clear()

                ref_point = []

            # END

        self.release()
        print('{} quit!'.format(self.__class__.__name__))


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--config',
        type=str,
        default='default_params.json',
        help='SpireCV2 Config (.json)')
    parser.add_argument(
        '--job-name', '-j',
        type=str,
        default='live',
        help='SpireCV Job Name')
    parser.add_argument(
        '--ip',
        type=str,
        default='127.0.0.1',
        help='SpireMS Core IP')
    parser.add_argument(
        '--port',
        type=int,
        default=9094,
        help='SpireMS Core Port')
    # args = parser.parse_args()
    args, unknown_args = parser.parse_known_args()
    if not os.path.isabs(args.config):
        current_path = os.path.abspath(__file__)
        params_dir = os.path.join(current_path[:current_path.find('spirecv-pro') + 11], 'params', 'spirecv2')
        args.config = os.path.join(params_dir, args.config)
    print("--config:", args.config)
    print("--job-name:", args.job_name)
    extra = get_extra_args(unknown_args)

    node = ClickTrackControlNode(args.job_name, param_dict_or_file=args.config, ip=args.ip, port=args.port, **extra)
    node.join()
